import re
from find_title_line import FindTitleLine
from title_number import TitleNumber


class EditTitleNumber:
    def __init__(self):

        self.r_sign_number = re.compile(r"(?P<sign>#+)\s+[\d.]*\s*")

    def add_number_one_line(self, line: str, level: int, number_s: str):
        new_line = self.r_sign_number.sub(
            lambda m: f"{m.group('sign')} {number_s} ", line, count=1)
        return new_line

    def remove_number_one_line(self, line: str):
        new_line = self.r_sign_number.sub(
            lambda m: f"{m.group('sign')} ", line, count=1)
        return new_line

    def add_number(self, lines):
        f = FindTitleLine()
        t = TitleNumber()
        index, levels = [], []
        for i, line in enumerate(lines):
            level = f(line)
            if level > 0:
                index.append(i)
                levels.append(level)
                t.append("#"*level)
        number_s_s = t.get_numbers()
        title_lines = [self.add_number_one_line(
            lines[i], level, number_s) for i, level, number_s in zip(index, levels, number_s_s)]
        for i, j in enumerate(index):
            lines[j] = title_lines[i]
        return lines

    def remove_number(self, lines):
        f = FindTitleLine()
        for i, line in enumerate(lines):
            level = f(line)
            if level > 0:
                lines[i] = self.remove_number_one_line(line)
        return lines


if __name__ == "__main__":
    # markdown 文档标题添加序号
    import os

    editor = EditTitleNumber()

    path = "example.md"
    lines, levels, number_s_s = [], [], []
    with open(path, encoding="utf8") as fr:
        lines = fr.readlines()

    name, ext = os.path.splitext(path)
    with open(name + "_with_title_number" + ext, 'w', encoding='utf8') as fw:
        fw.writelines(editor.add_number(lines))
        # fw.writelines(editor.remove_number(lines))
